# VOC数据集可视化
import os
import random

import numpy as np

from PIL import Image, ImageDraw, ImageFont

from tqdm import tqdm
import argparse
import xml.etree.ElementTree as ET

from utils.visualization import visual_color

import warnings
warnings.filterwarnings("ignore",category=DeprecationWarning)
 
def xml_reader(filename):
    """ Parse a PASCAL VOC xml file """
    tree = ET.parse(filename)
    objects = []
    for obj in tree.findall('object'):
        obj_struct = {}
        obj_struct['name'] = obj.find('name').text
        bbox = obj.find('bndbox')
        obj_struct['bbox'] = [int(bbox.find('xmin').text),
                              int(bbox.find('ymin').text),
                              int(bbox.find('xmax').text),
                              int(bbox.find('ymax').text)]
        objects.append(obj_struct)
 
    return objects
def get_image_list1(dataset_path, img_path, dataset_split):
    img_list = []
    dataset_split_path = os.path.join(dataset_path,'ImageSets/Main', dataset_split+'.txt')
    img_names = open(dataset_split_path, 'r').readlines()
    
    for img_name in img_names:
        img_list.append(os.path.join(dataset_path, img_path, img_name.strip() + '.jpg'))
    return img_list
def get_image_list(image_dir, suffix=['jpg','png']):
    '''get all image path ends with suffix'''
    if not os.path.exists(image_dir):
        print("PATH:%s not exists" % image_dir)
        return []
    imglist = []
    for root, sdirs, files in os.walk(image_dir):
        if not files:
            continue
        for filename in files:
            filepath = os.path.join(root, filename)
            if filename.split('.')[-1] in suffix:
                imglist.append(filepath)
    return imglist
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='VOC dataset visualization')
    parser.add_argument('--dataset', type=str, default='NWPUv2', help='Dataset:')
    parser.add_argument('--dataset_split',  default='test', type=str)
    parser.add_argument('--visual_path', default='exps/visualization/dataset_visual', help='The output dir of images', type=str)
    parser.add_argument('--num', default=None, help='The number of images you want to check', type=int)
    args = parser.parse_args()
    
    args.visual_path = os.path.join(args.visual_path, args.dataset, args.dataset_split)

    if not os.path.exists(args.visual_path):
        os.makedirs(args.visual_path)

    if args.dataset == 'NWPUv1':
        dataset_path = 'dataset/NWPUv1'
        classes_path ='model_data/nwpuv1/nwpuv1_classes.txt'
        img_path = 'JPEGImages'
        ann_path = 'Annotations'

    elif args.dataset == 'NWPUv2':
        dataset_path = 'dataset/NWPUv2'
        classes_path ='model_data/nwpuv2/nwpuv2_classes.txt'
        img_path = 'JPEGImages'
        ann_path = 'Annotations'

    elif args.dataset == 'DIOR':
        dataset_path = 'dataset/DIOR'
        classes_path = 'model_data/dior/dior_classes.txt'
        img_path = 'JPEGImages'
        ann_path = 'Annotations/Horizontal Bounding Boxes'

    else:
        print("Error: unrecognized dataset")
    
    img_dir_path = os.path.join(dataset_path, img_path)

    colors = visual_color(classes_path)
    
    # img_list = get_image_list(img_dir_path)
    img_list = get_image_list1(dataset_path, img_path, args.dataset_split)

    if args.num is not None:
        img_list = random.sample(img_list, args.num)
 
    for img_path in tqdm(img_list):
        img_name = img_path.lstrip(img_dir_path)
        image = Image.open(img_path)
        if image is None:
            continue
        xml_path = img_path.replace("JPEGImages", ann_path).replace(".jpg", ".xml").replace(".png", ".xml")
        objects = xml_reader(xml_path)
        if len(objects) == 0:
            continue

        font        = ImageFont.truetype(font='model_data/simhei.ttf', size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        thickness   = int(max((image.size[0] + image.size[1]) // np.mean(image.size), 1))

        # draw box and name
        for obj in objects:
            name = obj['name']
            box = obj['bbox']

            left, top, right, bottom  = box

            top     = max(0, np.floor(top).astype('int32'))
            left    = max(0, np.floor(left).astype('int32'))
            bottom  = min(image.size[1], np.floor(bottom).astype('int32'))
            right   = min(image.size[0], np.floor(right).astype('int32'))
        
            draw = ImageDraw.Draw(image)
            label = '{}'.format(name)
            label_size = draw.textsize(label, font)
            label = label.encode('utf-8')

            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])


            for i in range(thickness):
                draw.rectangle([left + i, top + i, right - i, bottom - i], outline=colors[name])
            draw.rectangle([tuple(text_origin), tuple(text_origin + label_size)], fill=colors[name])
            draw.text(text_origin, str(label,'UTF-8'), fill=(0, 0, 0), font=font)

            del draw

        # image.save(os.path.join(args.visual_path, img_name.replace(".jpg", ".png")), quality=95, subsampling=0)
        image.save(os.path.join(args.visual_path, img_name))
